{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5.3、5.4  PyTorchでDQN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# パッケージのimport\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import gym\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 動画の描画関数の宣言\n",
    "# 参考URL http://nbviewer.jupyter.org/github/patrickmineault\n",
    "# /xcorr-notebooks/blob/master/Render%20OpenAI%20gym%20as%20GIF.ipynb\n",
    "from JSAnimation.IPython_display import display_animation\n",
    "from matplotlib import animation\n",
    "from IPython.display import display\n",
    "\n",
    "\n",
    "def display_frames_as_gif(frames):\n",
    "    \"\"\"\n",
    "    Displays a list of frames as a gif, with controls\n",
    "    \"\"\"\n",
    "    plt.figure(figsize=(frames[0].shape[1]/72.0, frames[0].shape[0]/72.0),\n",
    "               dpi=72)\n",
    "    patch = plt.imshow(frames[0])\n",
    "    plt.axis('off')\n",
    "\n",
    "    def animate(i):\n",
    "        patch.set_data(frames[i])\n",
    "\n",
    "    anim = animation.FuncAnimation(plt.gcf(), animate, frames=len(frames),\n",
    "                                   interval=50)\n",
    "\n",
    "    anim.save('movie_cartpole_DQN.mp4')  # 動画のファイル名と保存です\n",
    "    display(display_animation(anim, default_mode='loop'))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 本コードでは、namedtupleを使用します。\n",
    "# namedtupleを使うことで、値をフィールド名とペアで格納できます。\n",
    "# すると値に対して、フィールド名でアクセスできて便利です。\n",
    "# https://docs.python.jp/3/library/collections.html#collections.namedtuple\n",
    "# 以下は使用例です\n",
    "\n",
    "from collections import namedtuple\n",
    "\n",
    "Tr = namedtuple('tr', ('name_a', 'value_b'))\n",
    "Tr_object = Tr('名前Aです', 100)\n",
    "\n",
    "print(Tr_object)  # 出力：tr(name_a='名前Aです', value_b=100)\n",
    "print(Tr_object.value_b)  # 出力：100\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# namedtupleを生成\n",
    "from collections import namedtuple\n",
    "\n",
    "Transition = namedtuple(\n",
    "    'Transition', ('state', 'action', 'next_state', 'reward'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定数の設定\n",
    "ENV = 'CartPole-v1'  # 使用する課題名\n",
    "GAMMA = 0.99  # 時間割引率\n",
    "MAX_STEPS = 200  # 1試行のstep数\n",
    "NUM_EPISODES = 500  # 最大試行回数\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 経験を保存するメモリクラスを定義します\n",
    "\n",
    "\n",
    "class ReplayMemory:\n",
    "\n",
    "    def __init__(self, CAPACITY):\n",
    "        self.capacity = CAPACITY  # メモリの最大長さ\n",
    "        self.memory = []  # 経験を保存する変数\n",
    "        self.index = 0  # 保存するindexを示す変数\n",
    "\n",
    "    def push(self, state, action, state_next, reward):\n",
    "        '''transition = (state, action, state_next, reward)をメモリに保存する'''\n",
    "\n",
    "        if len(self.memory) < self.capacity:\n",
    "            self.memory.append(None)  # メモリが満タンでないときは足す\n",
    "\n",
    "        # namedtupleのTransitionを使用し、値とフィールド名をペアにして保存します\n",
    "        self.memory[self.index] = Transition(state, action, state_next, reward)\n",
    "\n",
    "        self.index = (self.index + 1) % self.capacity  # 保存するindexを1つずらす\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        '''batch_size分だけ、ランダムに保存内容を取り出す'''\n",
    "        return random.sample(self.memory, batch_size)\n",
    "\n",
    "    def __len__(self):\n",
    "        '''関数lenに対して、現在の変数memoryの長さを返す'''\n",
    "        return len(self.memory)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# エージェントが持つ脳となるクラスです、DQNを実行します\n",
    "# Q関数をディープラーニングのネットワークをクラスとして定義\n",
    "\n",
    "import random\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "# import warnings\n",
    "# 忽略特定的警告信息\n",
    "# warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
    "# warnings.filterwarnings(\"ignore\")\n",
    "BATCH_SIZE = 32\n",
    "CAPACITY = 10000\n",
    "\n",
    "\n",
    "class Brain:\n",
    "    def __init__(self, num_states, num_actions):\n",
    "        self.num_actions = num_actions  # CartPoleの行動（右に左に押す）の2を取得\n",
    "\n",
    "        # 経験を記憶するメモリオブジェクトを生成\n",
    "        self.memory = ReplayMemory(CAPACITY)\n",
    "\n",
    "        # ニューラルネットワークを構築\n",
    "        self.model = nn.Sequential()\n",
    "        self.model.add_module('fc1', nn.Linear(num_states, 32))\n",
    "        self.model.add_module('relu1', nn.ReLU())\n",
    "        self.model.add_module('fc2', nn.Linear(32, 32))\n",
    "        self.model.add_module('relu2', nn.ReLU())\n",
    "        self.model.add_module('fc3', nn.Linear(32, num_actions))\n",
    "\n",
    "        print(self.model)  # ネットワークの形を出力\n",
    "\n",
    "        # 最適化手法の設定\n",
    "        self.optimizer = optim.Adam(self.model.parameters(), lr=0.0001)\n",
    "\n",
    "    def replay(self):\n",
    "        '''Experience Replayでネットワークの結合パラメータを学習'''\n",
    "\n",
    "        # -----------------------------------------\n",
    "        # 1. メモリサイズの確認\n",
    "        # -----------------------------------------\n",
    "        # 1.1 メモリサイズがミニバッチより小さい間は何もしない\n",
    "        if len(self.memory) < BATCH_SIZE:\n",
    "            return\n",
    "\n",
    "        # -----------------------------------------\n",
    "        # 2. ミニバッチの作成\n",
    "        # -----------------------------------------\n",
    "        # 2.1 メモリからミニバッチ分のデータを取り出す\n",
    "        transitions = self.memory.sample(BATCH_SIZE)\n",
    "\n",
    "        # 2.2 各変数をミニバッチに対応する形に変形\n",
    "        # transitionsは1stepごとの(state, action, state_next, reward)が、BATCH_SIZE分格納されている\n",
    "        # つまり、(state, action, state_next, reward)×BATCH_SIZE\n",
    "        # これをミニバッチにしたい。つまり\n",
    "        # (state×BATCH_SIZE, action×BATCH_SIZE, state_next×BATCH_SIZE, reward×BATCH_SIZE)にする\n",
    "        batch = Transition(*zip(*transitions))\n",
    "\n",
    "        # 2.3 各変数の要素をミニバッチに対応する形に変形し、ネットワークで扱えるようVariableにする\n",
    "        # 例えばstateの場合、[torch.FloatTensor of size 1x4]がBATCH_SIZE分並んでいるのですが、\n",
    "        # それを torch.FloatTensor of size BATCH_SIZEx4 に変換します\n",
    "        # 状態、行動、報酬、non_finalの状態のミニバッチのVariableを作成\n",
    "        # catはConcatenates（結合）のことです。\n",
    "        state_batch = torch.cat(batch.state)\n",
    "        action_batch = torch.cat(batch.action)\n",
    "        reward_batch = torch.cat(batch.reward)\n",
    "        non_final_next_states = torch.cat([s for s in batch.next_state\n",
    "                                           if s is not None])\n",
    "\n",
    "        # -----------------------------------------\n",
    "        # 3. 教師信号となるQ(s_t, a_t)値を求める\n",
    "        # -----------------------------------------\n",
    "        # 3.1 ネットワークを推論モードに切り替える\n",
    "        self.model.eval()\n",
    "\n",
    "        # 3.2 ネットワークが出力したQ(s_t, a_t)を求める\n",
    "        # self.model(state_batch)は、右左の両方のQ値を出力しており\n",
    "        # [torch.FloatTensor of size BATCH_SIZEx2]になっている。\n",
    "        # ここから実行したアクションa_tに対応するQ値を求めるため、action_batchで行った行動a_tが右か左かのindexを求め\n",
    "        # それに対応するQ値をgatherでひっぱり出す。\n",
    "        state_action_values = self.model(state_batch).gather(1, action_batch)\n",
    "\n",
    "        # 3.3 max{Q(s_t+1, a)}値を求める。ただし次の状態があるかに注意。\n",
    "\n",
    "        # cartpoleがdoneになっておらず、next_stateがあるかをチェックするインデックスマスクを作成\n",
    "        non_final_mask = torch.ByteTensor(tuple(map(lambda s: s is not None,\n",
    "                                                    batch.next_state)))\n",
    "        # まずは全部0にしておく\n",
    "        next_state_values = torch.zeros(BATCH_SIZE)\n",
    "\n",
    "        # 次の状態があるindexの最大Q値を求める\n",
    "        # 出力にアクセスし、max(1)で列方向の最大値の[値、index]を求めます\n",
    "        # そしてそのQ値（index=0）を出力します\n",
    "        # detachでその値を取り出します\n",
    "        next_state_values[non_final_mask] = self.model(\n",
    "            non_final_next_states).max(1)[0].detach()\n",
    "\n",
    "        # 3.4 教師となるQ(s_t, a_t)値を、Q学習の式から求める\n",
    "        expected_state_action_values = reward_batch + GAMMA * next_state_values\n",
    "\n",
    "        # -----------------------------------------\n",
    "        # 4. 結合パラメータの更新\n",
    "        # -----------------------------------------\n",
    "        # 4.1 ネットワークを訓練モードに切り替える\n",
    "        self.model.train()\n",
    "\n",
    "        # 4.2 損失関数を計算する（smooth_l1_lossはHuberloss）\n",
    "        # expected_state_action_valuesは\n",
    "        # sizeが[minbatch]になっているので、unsqueezeで[minibatch x 1]へ\n",
    "        loss = F.smooth_l1_loss(state_action_values,\n",
    "                                expected_state_action_values.unsqueeze(1))\n",
    "\n",
    "        # 4.3 結合パラメータを更新する\n",
    "        self.optimizer.zero_grad()  # 勾配をリセット\n",
    "        loss.backward()  # バックプロパゲーションを計算\n",
    "        self.optimizer.step()  # 結合パラメータを更新\n",
    "\n",
    "    def decide_action(self, state, episode):\n",
    "        '''現在の状態に応じて、行動を決定する'''\n",
    "        # ε-greedy法で徐々に最適行動のみを採用する\n",
    "        epsilon = 0.5 * (1 / (episode + 1))\n",
    "\n",
    "        if epsilon <= np.random.uniform(0, 1):\n",
    "            self.model.eval()  # ネットワークを推論モードに切り替える\n",
    "            with torch.no_grad():\n",
    "                action = self.model(state).max(1)[1].view(1, 1)\n",
    "            # ネットワークの出力の最大値のindexを取り出します = max(1)[1]\n",
    "            # .view(1,1)は[torch.LongTensor of size 1]　を size 1x1 に変換します\n",
    "\n",
    "        else:\n",
    "            # 0,1の行動をランダムに返す\n",
    "            action = torch.LongTensor(\n",
    "                [[random.randrange(self.num_actions)]])  # 0,1の行動をランダムに返す\n",
    "            # actionは[torch.LongTensor of size 1x1]の形になります\n",
    "\n",
    "        return action\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CartPoleで動くエージェントクラスです、棒付き台車そのものになります\n",
    "\n",
    "\n",
    "class Agent:\n",
    "    def __init__(self, num_states, num_actions):\n",
    "        '''課題の状態と行動の数を設定する'''\n",
    "        self.brain = Brain(num_states, num_actions)  # エージェントが行動を決定するための頭脳を生成\n",
    "\n",
    "    def update_q_function(self):\n",
    "        '''Q関数を更新する'''\n",
    "        self.brain.replay()\n",
    "\n",
    "    def get_action(self, state, episode):\n",
    "        '''行動を決定する'''\n",
    "        action = self.brain.decide_action(state, episode)\n",
    "        return action\n",
    "\n",
    "    def memorize(self, state, action, state_next, reward):\n",
    "        '''memoryオブジェクトに、state, action, state_next, rewardの内容を保存する'''\n",
    "        self.brain.memory.push(state, action, state_next, reward)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CartPoleを実行する環境のクラスです\n",
    "\n",
    "\n",
    "class Environment:\n",
    "\n",
    "    def __init__(self):\n",
    "        self.env = gym.make(ENV,render_mode='rgb_array')  # 実行する課題を設定\n",
    "        num_states = self.env.observation_space.shape[0]  # 課題の状態数4を取得\n",
    "        num_actions = self.env.action_space.n  # CartPoleの行動（右に左に押す）の2を取得\n",
    "        self.agent = Agent(num_states, num_actions)  # 環境内で行動するAgentを生成\n",
    "\n",
    "        \n",
    "    def run(self):\n",
    "        '''実行'''\n",
    "        episode_10_list = np.zeros(10)  # 10試行分の立ち続けたstep数を格納し、平均ステップ数を出力に利用\n",
    "        complete_episodes = 0  # 195step以上連続で立ち続けた試行数\n",
    "        episode_final = False  # 最後の試行フラグ\n",
    "        frames = []  # 最後の試行を動画にするために画像を格納する変数\n",
    "\n",
    "        for episode in range(NUM_EPISODES):  # 最大試行数分繰り返す\n",
    "            observation = self.env.reset()[0]  # 環境の初期化\n",
    "\n",
    "            state = observation  # 観測をそのまま状態sとして使用\n",
    "            state = torch.from_numpy(state).type(\n",
    "                torch.FloatTensor)  # NumPy変数をPyTorchのテンソルに変換\n",
    "            state = torch.unsqueeze(state, 0)  # size 4をsize 1x4に変換\n",
    "\n",
    "            for step in range(MAX_STEPS):  # 1エピソードのループ\n",
    "\n",
    "                if episode_final is True:  # 最終試行ではframesに各時刻の画像を追加していく\n",
    "                    frames.append(self.env.render())\n",
    "\n",
    "                action = self.agent.get_action(state, episode)  # 行動を求める\n",
    "\n",
    "                # 行動a_tの実行により、s_{t+1}とdoneフラグを求める\n",
    "                # actionから.item()を指定して、中身を取り出す\n",
    "                observation_next, _, done, _, _ = self.env.step(\n",
    "                    action.item())  # rewardとinfoは使わないので_にする\n",
    "\n",
    "                # 報酬を与える。さらにepisodeの終了評価と、state_nextを設定する\n",
    "                if done:  # ステップ数が200経過するか、一定角度以上傾くとdoneはtrueになる\n",
    "                    state_next = None  # 次の状態はないので、Noneを格納\n",
    "\n",
    "                    # 直近10episodeの立てたstep数リストに追加\n",
    "                    episode_10_list = np.hstack(\n",
    "                        (episode_10_list[1:], step + 1))\n",
    "\n",
    "                    if step < 195:\n",
    "                        reward = torch.FloatTensor(\n",
    "                            [-1.0])  # 途中でこけたら罰則として報酬-1を与える\n",
    "                        complete_episodes = 0  # 連続成功記録をリセット\n",
    "                    else:\n",
    "                        reward = torch.FloatTensor([1.0])  # 立ったまま終了時は報酬1を与える\n",
    "                        complete_episodes = complete_episodes + 1  # 連続記録を更新\n",
    "                else:\n",
    "                    reward = torch.FloatTensor([0.0])  # 普段は報酬0\n",
    "                    state_next = observation_next  # 観測をそのまま状態とする\n",
    "                    state_next = torch.from_numpy(state_next).type(\n",
    "                        torch.FloatTensor)  # numpy変数をPyTorchのテンソルに変換\n",
    "                    state_next = torch.unsqueeze(state_next, 0)  # size 4をsize 1x4に変換\n",
    "\n",
    "                # メモリに経験を追加\n",
    "                self.agent.memorize(state, action, state_next, reward)\n",
    "\n",
    "                # Experience ReplayでQ関数を更新する\n",
    "                self.agent.update_q_function()\n",
    "\n",
    "                # 観測の更新\n",
    "                state = state_next\n",
    "\n",
    "                # 終了時の処理\n",
    "                if done:\n",
    "                    print('%d Episode: Finished after %d steps：10次试验平均step数 = %.1lf' % (\n",
    "                        episode, step + 1, episode_10_list.mean()))\n",
    "                    break\n",
    "\n",
    "            if episode_final is True:\n",
    "                # 動画を保存と描画\n",
    "                display_frames_as_gif(frames)\n",
    "                break\n",
    "\n",
    "            # 10連続で200step経ち続けたら成功\n",
    "            if complete_episodes >= 10:\n",
    "                print('10回连续成功')\n",
    "                episode_final = True  # 次の試行を描画を行う最終試行とする\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# main クラス\n",
    "cartpole_env = Environment()\n",
    "cartpole_env.run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
